"""用于生成每个陨石坑对应的标注文件"""

import os
import cv2
import numpy as np
import math

LATITUDE_RANGE = 120  # degree
LONGITUDE_RANGE = 360  # degree
X_PIXEL_RANGE = 61440
Y_PIXEL_RANGE = 184320
MOON_RADIUS = 1737_400  # m
# METER_PER_PIXEL = 59.225
BATCH_SIZE = 960


def sgn(x):
    if x > 0:
        return 1
    elif x < 0:
        return -1
    else:
        return 0


# Here is the definition of image coordinate system
# +-------> y
# | (0, 0)
# |
# |
# |
# v
# x
def mercator_projection(lat, lon):
    """
    Warning ! The unit of latitude and longitude is degree.
    """
    # 这里要加上一半的图像宽度，这是因为标注数据集中的角度是从-60度到+60度，起始位置不为零，因此需要修正。
    x = (60 - lat) * X_PIXEL_RANGE / LATITUDE_RANGE
    # 这里不加一半的图像高度，这是因为标注数据集中的角度是从0到360度的，但是，图像的经度是从-180到+180度的，因此需要修正。
    if lon > 180:
        # y = (lon - 180) * X_PIXEL_RANGE * math.cos(math.radians(lat)) / LATITUDE_RANGE
        y = (lon - 180) * Y_PIXEL_RANGE / LONGITUDE_RANGE
    else:
        # y = (lon + 180) * X_PIXEL_RANGE * math.cos(math.radians(lat)) / LATITUDE_RANGE
        y = (lon + 180) * Y_PIXEL_RANGE / LONGITUDE_RANGE
    return x, y


def labels(line, header):
    # for circle projection
    temp_1, temp_2 = mercator_projection(float(line[1]), float(line[2]))
    x_folder = int(temp_1) // BATCH_SIZE
    y_folder = int(temp_2) // BATCH_SIZE
    x = temp_1 - BATCH_SIZE * x_folder
    y = temp_2 - BATCH_SIZE * y_folder
    cir_labels = {"Y_FOLDER": y_folder, "X_FOLDER": x_folder, "Y_SUB": y, "X_SUB": x}
    # for ellipse projection
    temp_1, temp_2 = mercator_projection(float(line[3]), float(line[4]))
    x_folder = int(temp_1) // BATCH_SIZE
    y_folder = int(temp_2) // BATCH_SIZE
    x = temp_1 - BATCH_SIZE * x_folder
    y = temp_2 - BATCH_SIZE * y_folder
    elli_labels = {"Y_FOLDER": y_folder, "X_FOLDER": x_folder, "Y_SUB": y, "X_SUB": x}

    for field, val in zip(header, line):
        if "CIR" in field:
            cir_labels = dict(**{field: val}, **cir_labels)
        elif "ELLI" in field:
            elli_labels = dict(**{field: val}, **elli_labels)
        else:
            cir_labels = dict(**{field: val}, **cir_labels)
            elli_labels = dict(**{field: val}, **elli_labels)
    return cir_labels, elli_labels


if __name__ == "__main__":
    # 从标注文件中获取某一行
    LABEL_DATA_PATH = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater_database_robbins_2018_bundle/data/lunar_crater_database_robbins_2018.csv"

    SUBFIG_DIR = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"
    with open(LABEL_DATA_PATH, "r") as f:
        # 有部分行只有圆的数据，没有椭圆的数据
        # 流式读取，第一行是标题，并使用标题的字段名作为数据变量名
        header = f.readline()
        header = header.strip().split(",")
        while True:
            # 读取第一行数据
            line_ = f.readline()
            if not line_:
                break
            line = line_.strip().split(",")
            try:
                cir_labels, elli_labels = labels(line, header)
            except:
                continue
            cir_x_folder = cir_labels["X_FOLDER"]
            cir_y_folder = cir_labels["Y_FOLDER"]
            elli_x_folder = elli_labels["X_FOLDER"]
            elli_y_folder = elli_labels["Y_FOLDER"]
            if (
                cir_x_folder > 64
                or cir_x_folder < 0
                or elli_x_folder > 64
                or elli_x_folder < 0
            ):
                continue
            if not os.path.exists(f"{SUBFIG_DIR}/labels/craters"):
                os.makedirs(f"{SUBFIG_DIR}/labels/craters")
            img = cv2.imread(
                f"{SUBFIG_DIR}/textures/{cir_x_folder}/txr_{cir_x_folder}_{cir_y_folder}_{BATCH_SIZE}.png",
                cv2.IMREAD_GRAYSCALE,
            )
            radius = float(cir_labels["DIAM_CIRC_IMG"]) / 2
            x = int(cir_labels["X_SUB"])
            y = int(cir_labels["Y_SUB"])
            R_x = int(
                math.degrees(radius / MOON_RADIUS) * X_PIXEL_RANGE / LATITUDE_RANGE
            )
            R_y = int(
                math.degrees(radius / MOON_RADIUS)
                * Y_PIXEL_RANGE
                / LONGITUDE_RANGE
                / math.cos(math.radians(float(cir_labels["LAT_CIRC_IMG"])))
            )
            if radius > 20:
                continue
            else:
                if x - radius < 0 or x + radius > 960:
                    continue
                elif y - radius < 0 or y + radius > 960:
                    continue
                else:
                    sub_img = img[y - R_y : y + R_y, x - R_x : x + R_x]
                    cv2.imwrite(
                        f"{SUBFIG_DIR}/labels/craters/{cir_labels['CRATER_ID']}.png",
                        sub_img,
                    )
